package com.jboltai.scheduler;


import com.alibaba.ttl.threadpool.TtlExecutors;
import com.jboltai.config.JBoltAIConfig;
import com.jboltai.event.Event;
import com.jboltai.event.EventError;
import com.jboltai.event.EventErrorType;
import com.jboltai.event.EventState;
import com.jboltai.event.ai.AIEvent;
import com.jboltai.event.ai.aitext.AITextEvent;
import com.jboltai.event.embedding.EmbeddingEvent;
import com.jboltai.event.signal.FailEvent;
import com.jboltai.event.signal.SuccessEvent;
import com.jboltai.event.text.TextEvent;
import com.jboltai.util.other.ExpiringCache;
import com.jboltai.util.share.AIEventThreadShareData;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * 事件总线
 */
public class EventBus {
    private static final Logger LOGGER = LoggerFactory.getLogger(JBoltAIConfig.LOGGER_NAME);
    /**
     * ai事件调度器
     */
    static AIEventScheduler aiEventScheduler = new AIEventScheduler();
    static TextEventScheduler textEventScheduler = new TextEventScheduler();

    static EmbeddingEventScheduler embeddingEventScheduler = new EmbeddingEventScheduler();

    static AITextEventScheduler aiTextEventScheduler = new AITextEventScheduler();


    static ScheduledExecutorService executorService = TtlExecutors.getTtlScheduledExecutorService(Executors.newScheduledThreadPool(2));
    static {
        executorService.scheduleWithFixedDelay(aiEventScheduler,0,5, TimeUnit.MILLISECONDS);
        executorService.scheduleWithFixedDelay(textEventScheduler,0,5, TimeUnit.MILLISECONDS);
        executorService.scheduleWithFixedDelay(embeddingEventScheduler,0,5, TimeUnit.MILLISECONDS);
    }

    /**
     * 将事件id和事件对象 进行映射，方便根据id获取事件
     */
    private static final ConcurrentHashMap<String, Event> EVENT_HOLDER = new ConcurrentHashMap<>();

    /**
     * 缓存要取消的key
     */
    private static final ExpiringCache<String, String> cancelKeyCache = new ExpiringCache<>(30000);

    /**
     * 发布事件，将事件插入队列等待执行
     * @param event
     */
    public static <T extends Event> T publish(T event) {
        if (event == null) {
            throw new NullPointerException("事件不可为null");
        }
        if (event.getId() != null && cancelKeyCache.containsKey(event.getId())) {
            try{
                //该事件已被取消
                event.cancel();
                return event;
            } finally {
                cancelKeyCache.remove(event.getId());
            }
        }

        String preCheckMsg = event.preCheck();
        if (preCheckMsg != null) {
            event.fail(EventErrorType.PREV_CHECK_NOT_PASS, preCheckMsg);
            return event;
        }
        AIEventThreadShareData threadShareData ;
        if (event.getThreadShareData() == null && (threadShareData = AIEventThreadShareData.getShareData()) != null) {
            event.setThreadShareData(threadShareData);
        }
        switch (event.getType()) {
            case AI_CHAT:
                AIEvent aiEvent = (AIEvent) event;
                aiEventScheduler.publish(aiEvent);
                break;

            case AI_KNOWLEDGE:
            case AI_FUNC_CALL:
            case AI_TEXT_TO_SQL:
            case AI_TEXT_TO_JSON:
            case AI_INTENTION:
            case AI_REWRITE_PROMPT:

                aiTextEventScheduler.publish((AITextEvent) event);
                break;
            case TEXT_EXTRACT:
            case TEXT_SPLIT:
                textEventScheduler.publish((TextEvent) event);
                break;
            case EMBEDDING:
                embeddingEventScheduler.publish((EmbeddingEvent) event);
                break;
            case SIGNAL_SUCCESS:{
                event.setState(EventState.PROCESSING);
                event.success(((SuccessEvent)event).getResult());
                break;
            }

            case SIGNAL_FAIL: {
                event.setState(EventState.PROCESSING);
                EventError error = ((FailEvent) event).getError();
                event.fail(error.getErrorType(), error.getMsg(), error.getException());
                break;
            }

        }

        hold(event);

        return event;

    }



    /**
     * 发布事件，并立即分配资源执行
     * @param event
     */
    public static <T extends Event> T dispatch(T event) {
        if (event == null) {
            throw new NullPointerException("事件不可为null");
        }
        if (event.getId() != null && cancelKeyCache.containsKey(event.getId())) {
            try{
                //该事件已被取消
                event.cancel();
                return event;
            } finally {
                cancelKeyCache.remove(event.getId());
            }
        }
        event.setState(EventState.WAITING);
        String preCheckMsg = event.preCheck();
        if (preCheckMsg != null) {
            event.fail(EventErrorType.PREV_CHECK_NOT_PASS, preCheckMsg);
            return event;
        }

        AIEventThreadShareData threadShareData ;
        if (event.getThreadShareData() == null && (threadShareData = AIEventThreadShareData.getShareData()) != null) {
            event.setThreadShareData(threadShareData);
        }
        //是要立马执行的
        event.setState(EventState.PROCESSING);
        switch (event.getType()) {
            case AI_CHAT:

                AIEvent aiEvent = (AIEvent) event;
                aiEventScheduler.dispatch(aiEvent);
                break;
            case AI_KNOWLEDGE:
            case AI_FUNC_CALL:
            case AI_TEXT_TO_SQL:
            case AI_TEXT_TO_JSON:
            case AI_INTENTION:
            case AI_REWRITE_PROMPT:
                aiTextEventScheduler.dispatch((AITextEvent) event);
                break;
            case TEXT_EXTRACT:
            case TEXT_SPLIT:
                textEventScheduler.dispatch((TextEvent) event);
                break;
            case EMBEDDING:
                embeddingEventScheduler.dispatch((EmbeddingEvent) event);
                break;

            case SIGNAL_SUCCESS:
                event.setState(EventState.PROCESSING);
                event.success(((SuccessEvent)event).getResult());
                break;
            case SIGNAL_FAIL: {
                EventError error = ((FailEvent) event).getError();
                event.fail(error.getErrorType(), error.getMsg(), error.getException());
                break;
            }
        }
        hold(event);
        return event;

    }


    /**
     * 取消某个事件的执行
     * @param id
     */
    public static void cancel(String id) {
        Event event = EVENT_HOLDER.get(id);
        if (event != null) {
            event.cancel();
        } else if (id != null){
            //这个事件还没创建，就先缓存起来
            cancelKeyCache.put(id, id);
        }
    }

    /**
     * 将事件纳入全局上下文
     * @param event
     */
    public static void hold(Event event) {
        if (event.getId() != null && !event.isSuccess() && !event.isFail()) {
            EVENT_HOLDER.put(event.getId(), event);
        }
    }

    /**
     * 事件执行完成后，将它从EVENT_HOLDER中移除
     * @param id
     */
    public static void unHold(String id) {
        if (id != null) {
            EVENT_HOLDER.remove(id);
        }
    }
    /**
     * 根据id获取event
     * @param id
     */
    public static Event get(String id) {
        if (id != null) {
            EVENT_HOLDER.remove(id);
        }
        return null;
    }

    /**
     * 根据id 检测是否存在
     * @param id
     */
    public static boolean exists(String id) {
        if (id != null) {
            return EVENT_HOLDER.containsKey(id);
        }
        return false;
    }


}
